package examples;

import java.io.File;
import java.util.Date;
import java.util.List;
import luceneqe.modelsIR.BM25MultiTermQuery;
import luceneqe.modelsIR.SimilarityBM25;
import org.apache.lucene.analysis.snowball.SnowballAnalyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.queryParser.QueryParser;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.TopDocCollector;
import org.apache.lucene.search.TopDocs;

import luceneqe.queryExpansion.modelsQE.QueryExpansionModel;
import luceneqe.queryExpansion.modelsQE.QueryExpansionModelFactory;
import luceneqe.utils.BM25Parameters;
import luceneqe.utils.CollectionStats;
import luceneqe.utils.ExpansionParameters;

public class BM25QueryExpansionKLD {

	public static void main(String[] args) throws Exception {

		  String query = "(house)";
		  String index = "../LuceneQE/examples/luceneqe-test-index";
		  String field = "text";

		String fields[] = { field };
		CollectionStats.setAvgDocumentLength(598.0f);
		CollectionStats.setFreqTokenCollection(34906437.0);
		CollectionStats.setTotalDocumentsCollection(110245.0);
		CollectionStats.setFields(fields);		
		ExpansionParameters.setBetaRocchio(0.3);
		ExpansionParameters.setIDFthreshold(2.5);
		ExpansionParameters.setNumExpansionTerms(1);
		ExpansionParameters.setTopNdocuments(5);
		BM25Parameters.setB(0.67f);
		BM25Parameters.setK1(3.8f);
		long start = new Date().getTime();
		File indexDir = new File(index);
		if (!indexDir.exists() || !indexDir.isDirectory()) {
			throw new Exception(indexDir
			+ " does not exist or is not a directory.");
		}
		SnowballAnalyzer analyzer = new SnowballAnalyzer("English");
		IndexSearcher is = new IndexSearcher(index);
		IndexReader ir = is.getIndexReader();
		QueryParser parser = new QueryParser(field, analyzer);

		is.setSimilarity(new SimilarityBM25());

		System.out.println("COLLECTION: " + index);
		System.out.println(" size: " + is.maxDoc());
		System.out.println("SIM: " + is.getSimilarity().getClass());
		System.out.println("QUERY: <" + query + ">");

		// Expand Query:
		Query queryBM25 = new BM25MultiTermQuery(parser.parse(query),fields, query);
		QueryExpansionModelFactory.selectQueryExpansionModel(QueryExpansionModelFactory.KLD);
		QueryExpansionModel thes = QueryExpansionModelFactory.getQueryExpansionModel(ir, is, query, is.search(queryBM25), fields);

		// Get candidate list and construct expanded query:
		List<String> candidateList = thes.getPlainList();
		String queryExpanded = candidateList.toString();
		System.out.println(" Candidate List: " + queryExpanded);
		queryExpanded = queryExpanded.replaceAll("\\[", "");
		queryExpanded = queryExpanded.replaceAll("\\]", "");
		queryExpanded = queryExpanded.replaceAll("\\,", "");
		System.out.println(" Expanded Query : " + queryExpanded);
		System.out.println();

		// Searching
		Query queryBM25Expanded = new BM25MultiTermQuery(parser.parse(queryExpanded), fields, queryExpanded);
		TopDocCollector collector = new TopDocCollector(1000);
		is.search(queryBM25Expanded, collector);

		TopDocs results = collector.topDocs();
		ScoreDoc[] doc = results.scoreDocs;
		for (int n = 0; n < doc.length; n++) {
			int d = doc[n].doc;
			float s = doc[n].score;
			Document document = ir.document(d);
			System.out.println(n + ": ID=" + document.get("id") + "s=" + s);
			System.out.println(is.explain(queryBM25, d).toString());
		}

		is.close();
		ir.close();
		long end = new Date().getTime();
		System.out.println("Time: " + (end - start) / 1000 + " segundos");
	}
}
